package net;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.net.Socket;
import java.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import types.Message;

public class EncryptedConnectionHandler implements ConnectionHandler {

	private Socket mSocket;
	private MessageReceivedListener mMessageReceivedListener;
	ObjectOutputStream oos;
	
	   private Cipher inputCipher;
	   private Cipher outputCipher;
	
	public EncryptedConnectionHandler(Socket socket, MessageReceivedListener messageReceivedListener, String key) throws Exception {
		
		if(key == null || key.length() != 16)
			throw new Exception("Key length must be 16");
		
		SecretKey secretKey = new SecretKeySpec(key.getBytes(), "AES");

		mSocket = socket;
		mMessageReceivedListener = messageReceivedListener;
		
	      try
	      {
	    	  inputCipher = Cipher.getInstance("AES/CBC/PKCS5Padding", "SunJCE");
	         int size = inputCipher.getBlockSize();
	         byte[] tmp = new byte[size];
	         Arrays.fill(tmp, (byte)15);
	         IvParameterSpec iv = new IvParameterSpec(tmp);
	         inputCipher.init(Cipher.DECRYPT_MODE, secretKey, iv);
	         
	         outputCipher = Cipher.getInstance("AES/CBC/PKCS5Padding", "SunJCE");
	         size = outputCipher.getBlockSize();
	         tmp = new byte[size];
	         Arrays.fill(tmp, (byte)15);
	         iv = new IvParameterSpec(tmp);
	         outputCipher.init(Cipher.ENCRYPT_MODE, secretKey, iv);
	      }
	      catch(Exception e)
	      {
	         e.printStackTrace();
	         throw new IOException("Failed to init cipher: "+e.getMessage());
	      }
		
		try {
			OutputStream os = mSocket.getOutputStream();
			oos = new ObjectOutputStream(os);
		} catch (IOException e1) {
			// TODO Auto-generated catch block
			e1.printStackTrace();
		}
		
		Thread read = new Thread(new Runnable() {
			
			@Override
			public void run() {

				try {
					InputStream is = mSocket.getInputStream();
					ObjectInputStream ois = new ObjectInputStream(is);

					
					while(mSocket.isConnected()) {
						String m = (String)ois.readObject();
						
						if(mMessageReceivedListener != null)
							mMessageReceivedListener.messageReceived(fromString(m));
					}
					
				} catch (IOException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				} catch (ClassNotFoundException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
		});
		
		read.start();
	}
	
	public void sendMessage(Message message) throws IOException {
		
		String encrypted = toString(message);
		oos.writeObject(encrypted);
	}

	public void closeSocket() {
		try {
			mSocket.getInputStream().close();
			mSocket.getOutputStream().close();
			mSocket.close();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
	}
	
	public MessageReceivedListener getMessageReceivedListener() {
		return mMessageReceivedListener;
	}

	public void setMessageReceivedListener(MessageReceivedListener mMessageReceivedListener) {
		this.mMessageReceivedListener = mMessageReceivedListener;
	}
	
    /** Read the object from Base64 string. */
    private Message fromString( String s ) throws IOException ,
                                                        ClassNotFoundException {
        byte [] data = Base64Coder.decode( s );
        byte[] decryppted;
		try {
			decryppted = inputCipher.doFinal(data);
	        ObjectInputStream ois = new ObjectInputStream( 
                    new ByteArrayInputStream(  decryppted ) );
	        	Message m  = (Message)ois.readObject();
	        	ois.close();
	        	return m;
		} catch (IllegalBlockSizeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return null;
    }

    
    private String toString( Serializable o ) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ObjectOutputStream oos = new ObjectOutputStream( baos );
        oos.writeObject( o );
        oos.close();
        
        byte[] encrypted;
		try {
			encrypted = outputCipher.doFinal(baos.toByteArray());
			return new String( Base64Coder.encode( encrypted ) );
		} catch (IllegalBlockSizeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (BadPaddingException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return null;
    }
	
}